import torch
import cv2
import os
import glob
from torch.utils.data import Dataset
import random
import matplotlib.pyplot as plt
import numpy as np

# 类别标签字符串
NEGATIVE = "0"
POSITIVE = "1"


class DataLoader(Dataset):
    def __init__(self, root_dir, target_dim=(512, 512)):
        # 初始化，加载指定目录下的图片
        self.root_dir = root_dir  # 数据根目录
        self.target_dim = target_dim  # 目标维度 (height, width)

        # 存储图片路径和标签
        self.image_files = []
        self.class_ids = []
        self.label_map = {NEGATIVE: 0, POSITIVE: 1}

        # 遍历类别文件夹并加载图片
        for label, label_id in self.label_map.items():
            class_folder = os.path.join(root_dir, label)
            image_glob = os.path.join(class_folder, '*.png')
            images = glob.glob(image_glob)
            self.image_files.extend(images)
            self.class_ids.extend([label_id] * len(images))

    def flip_image(self, img, flip_type):
        # 根据指定的翻转类型翻转图片
        return cv2.flip(img, flip_type)

    def __getitem__(self, idx):
        # 根据索引获取单个样本
        img_path = self.image_files[idx]
        class_id = self.class_ids[idx]
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, self.target_dim)
        img = img.reshape(1, *self.target_dim)  # 重塑为 (C, H, W)
        img = img.astype(np.float32) / 255.0

        # 随机翻转图片
        flip_type = random.choice([-1, 0, 1])
        img = self.flip_image(img, flip_type)

        # 返回图片和标签
        return img, torch.tensor(class_id, dtype=torch.long)

    def __len__(self):
        # 返回数据集大小
        return len(self.image_files)


def show_image_label(img, label):
    # 显示图片和对应的标签
    img = img.squeeze(0)  # 去除单通道维度
    label_names = [NEGATIVE, POSITIVE]
    label_str = label_names[label.item()]

    fig, axs = plt.subplots(1, 2, figsize=(10, 5))
    axs[0].imshow(img, cmap='gray')
    axs[0].set_title("Image")
    axs[0].axis('on')

    axs[1].imshow(img, cmap='gray')
    axs[1].set_title(f"Label: {label_str}")
    axs[1].axis('on')

    plt.show()


if __name__ == "__main__":
    # 数据集路径
    dataset_path = r"C:\Users\yinjie\Desktop\archive\train"
    dataset = DataLoader(dataset_path)

    print(f"Dataset length: {len(dataset)}")

    # 创建数据加载器
    batch_size = 1000
    train_data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)

    # 测试数据加载和显示
    for img_batch, label_batch in train_data_loader:
        print(f"Image batch shape: {img_batch.shape}")
        print(f"Label batch shape: {label_batch.shape}")

        # 显示第一张图片的标签和图片
        first_label = label_batch[0]
        first_label_str = [NEGATIVE, POSITIVE][first_label.item()]
        print(f"First Label (Integer): {first_label.item()}")
        print(f"First Label Name: {first_label_str}")

        show_image_label(img_batch[0], first_label)
